import numpy as np
from sklearn import svm
from Translate import Translate
from Evaluate import Evaluate

def train_svm(isEvaluate = True, task = 0):
    evaluate = Evaluate()

    if task == 0 or task == 1:
        data0 = np.load('./data_task1/traindata_part0_pca20.npz')
        data1 = np.load('./data_task1/traindata_part1_pca20.npz')
        if not isEvaluate:
            traindata0 = data0['datas']
            traindata1 = data1['datas']
            trainlabel0 = data0['labels']
            trainlabel1 = data1['labels']
            predictlabel0 = np.zeros((6, len(traindata0), 6))
            predictlabel1 = np.zeros((6, len(traindata1), 6))
        else:
            traindata0 = data0['datas'][0:-7]
            traindata1 = data1['datas'][0:-7]
            trainlabel0 = data0['labels'][:,0:-7,:]
            trainlabel1 = data1['labels'][:,0:-7,:]
            predictlabel0 = np.zeros((6, len(traindata0), 6))
            predictlabel1 = np.zeros((6, len(traindata1), 6))

        models = {}

        for i in range(6):
            for j in range(6):
                # model = svm.SVR(C=14, gamma=0.007)
                model = svm.SVR(C=17, gamma=0.01)
                Y = trainlabel0[i,:,j]
                model.fit(traindata0, Y)
                models[(i, j, 0)] = model
                predictlabel0[i,:,j] = model.predict(traindata0)

                # model = svm.SVR(C=0.1, gamma=0.001)
                model = svm.SVR(C=0.1, gamma=0.001)
                Y = trainlabel1[i,:,j]
                model.fit(traindata1, Y)
                models[(i, j, 1)] = model
                predictlabel1[i,:,j] = model.predict(traindata1)

        MAPE_train_0 = evaluate.evaluate_test(predictlabel0, trainlabel0)
        MAPE_train_1 = evaluate.evaluate_test(predictlabel1, trainlabel1)
        print 'travel time MAPE_train: ' + str(MAPE_train_0) + ' ' + str(MAPE_train_1)


        if not isEvaluate:
            data0 = np.load('./data_task1/testdata_part0_pca20.npz')
            data1 = np.load('./data_task1/testdata_part1_pca20.npz')
            testdata0 = data0['tests']
            testdata1 = data1['tests']
            testlabel = np.zeros((6, 7, 12))
        else:
            testdata0 = data0['datas'][-7:]
            testdata1 = data1['datas'][-7:]
            testlabel = np.zeros((6, 7, 12))
            reslabel = np.concatenate((data0['labels'][:,-7:,:], data1['labels'][:,-7:,:]), axis=2)

        for i in range(6):
            for j in range(6):
                testlabel[i,:,j] = models[(i, j, 0)].predict(testdata0) * 0.94
                testlabel[i,:,j+6] = models[(i, j, 1)].predict(testdata1) * 0.93

        # print testlabel

        if not isEvaluate:
            trans = Translate()
            trans.translate(testlabel, path = './result_svm', task = 1)
        else:
            MAPE = evaluate.evaluate_test(testlabel, reslabel)
            print 'travel time MAPE: ' + str(MAPE)

        if task == 1:
            return testlabel

    if task == 0 or task == 2:
        data0 = np.load('./data_task2/traindata_part0_normalizing.npz')
        data1 = np.load('./data_task2/traindata_part1_normalizing.npz')
        if not isEvaluate:
            traindata0 = data0['datas']
            traindata1 = data1['datas']
            trainlabel0 = data0['labels']
            trainlabel1 = data1['labels']
            predictlabel0 = np.zeros((5, len(traindata0), 6))
            predictlabel1 = np.zeros((5, len(traindata1), 6))
        else:
            traindata0 = data0['datas'][0:-7]
            traindata1 = data1['datas'][0:-7]
            trainlabel0 = data0['labels'][:,0:-7,:]
            trainlabel1 = data1['labels'][:,0:-7,:]
            predictlabel0 = np.zeros((5, len(traindata0), 6))
            predictlabel1 = np.zeros((5, len(traindata1), 6))

        models = {}

        for i in range(5):
            for j in range(6):
                model = svm.SVR(kernel='sigmoid', C=9)
                # model = svm.SVR(C=3, gamma=0.003)
                Y = trainlabel0[i,:,j]
                model.fit(traindata0, Y)
                models[(i, j, 0)] = model
                predictlabel0[i,:,j] = model.predict(traindata0)

                model = svm.SVR(kernel='rbf', C=75, gamma=0.003)
                # model = svm.SVR(kernel='sigmoid', C=5)
                # model = svm.SVR(C=2, gamma=0.002)
                Y = trainlabel1[i,:,j]
                model.fit(traindata1, Y)
                models[(i, j, 1)] = model
                predictlabel1[i,:,j] = model.predict(traindata1)


        MAPE_train_0 = evaluate.evaluate_test(predictlabel0, trainlabel0)
        MAPE_train_1 = evaluate.evaluate_test(predictlabel1, trainlabel1)
        print 'traffic volume MAPE_train_0/1: ' + str(MAPE_train_0) + ' ' + str(MAPE_train_1)
        
        if not isEvaluate:
            data0 = np.load('./data_task2/testdata_part0_normalizing.npz')
            data1 = np.load('./data_task2/testdata_part1_normalizing.npz')
            testdata0 = data0['tests']
            testdata1 = data1['tests']
            testlabel = np.zeros((5, 7, 12))
        else:
            testdata0 = data0['datas'][-7:]
            testdata1 = data1['datas'][-7:]
            testlabel = np.zeros((5, 7, 12))
            reslabel = np.concatenate((data0['labels'][:,-7:,:], data1['labels'][:,-7:,:]), axis=2)

        for i in range(5):
            for j in range(6):
                testlabel[i,:,j] = models[(i, j, 0)].predict(testdata0) * 0.95
                testlabel[i,:,j+6] = models[(i, j, 1)].predict(testdata1) * 0.91

        # print testlabel

        if not isEvaluate:
            trans = Translate()
            trans.translate(testlabel, path = './result_svm', task = 2)
        else:
            MAPE = evaluate.evaluate_test(testlabel, reslabel)
            print 'traffic volume MAPE: ' + str(MAPE)

        if task == 2:
            return testlabel

if __name__ == '__main__':
    train_svm(isEvaluate = True, task = 1)
